#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Mon Nov  7 12:23:46 2016

@author: kevin
"""
from cv2 import line as drawLine
import numpy as np
import cv2
class genStrokes(object):
     def __init__(self, config, ):
		 
          self.num = config.batch_sz * config.M
          self.im_sz = config.img_sz
          self.im_ch = config.img_channels
          
          self.window_size = np.array((config.win_size, config.win_size))
          self.slideWindowHeight = config.win_size
          self.slideWindowWidth = config.win_size
          self.stroke_len_range = (int(self.slideWindowHeight*0.5),int(self.slideWindowHeight*1.2))
          self.threathhold = 0.99
          self.step_width = 5
          self.step_height = 5
          self.strokeColorBg = (0,0,0)       # black stroke
          self.strokeColorFg = (255,255,255) # white stroke
     
     
     def plotRandLine(self, img, loc, col):
          left_top = loc[0]
          right_down = loc[1]

          while True:
               start_loc_x = int(np.random.uniform(left_top[0],right_down[0]))
               start_loc_y = int(np.random.uniform(left_top[1],right_down[1]))
               start_loc = (start_loc_x,start_loc_y)
  
               end_loc_x = int(np.random.uniform(left_top[0],right_down[0]))
               end_loc_y = int(np.random.uniform(left_top[1],right_down[1]))
               end_loc = (end_loc_x,end_loc_y)
               
               length = np.sqrt(((np.array(end_loc)-np.array(start_loc))**2).sum())
               length = int(length)
               
               if length >= self.stroke_len_range[0] and length <= self.stroke_len_range[1]:
                    drawLine(img,start_loc,end_loc,col,3)
                    break
          return img
     
     def genStrokesForOne(self, alphaImg, rgbImg):
          img_height = alphaImg.shape[0]
          img_width = alphaImg.shape[1]
          i = j = 0
          while i <= img_height:
               while j <= img_width:
                    slideWin = alphaImg[i:i+self.slideWindowWidth,j:j+self.slideWindowHeight].copy()
                    bg_judge = (slideWin == 0).astype(np.int32)
                    num_pt_bg = bg_judge.sum().astype(np.float32)
                    num_pt_all = self.slideWindowHeight * self.slideWindowWidth
                    ratio = num_pt_bg/num_pt_all
                    if ratio > self.threathhold:          # background
                         rgbImg = self.plotRandLine(rgbImg,[(j,i),(j+self.slideWindowHeight,i+self.slideWindowWidth)],self.strokeColorBg)
                    elif ratio < (1 - self.threathhold):  # foreground
                         rgbImg = self.plotRandLine(rgbImg,[(j,i),(j+self.slideWindowHeight,i+self.slideWindowWidth)],self.strokeColorFg)
                    j+=int(self.slideWindowHeight*0.5+1)
               i+=int(self.slideWindowWidth*0.5+1)
          return rgbImg.copy()
     
     

     def genStroke(self, strokeMap, loc, rgbImg, alphaImg):
          leftup = [max(0,(loc[0] - self.window_size[0]/2)),max(0,(loc[1] - self.window_size[1]/2))]
          
          rightdown = [max(0,(loc[0] + self.window_size[0]/2)),max(0,(loc[1] + self.window_size[1]/2))]
          temp_alpha_img = alphaImg[leftup[1]:rightdown[1], leftup[0]:rightdown[0]]
          temp_rgb_img = rgbImg[leftup[1]:rightdown[1],leftup[0]:rightdown[0]].copy()
          
          strokes_img = self.genStrokesForOne(temp_alpha_img,temp_rgb_img)
          strokeMap[leftup[1]:rightdown[1],leftup[0]:rightdown[0]] = strokes_img
          return strokeMap
          
     
     def drawstroke(self, strokeMap, loc, image, gtAlpha):
          self.strokeMap = strokeMap
          self.loc = loc           # [N, 2]   N = batch*M
          self.image = image
          self.gtAlpha = gtAlpha
          
          for i in range(self.num):
               loc_this = self.loc[i,:]
               img_shape = np.array((self.im_sz, self.im_sz))
               loc_this = (loc_this*img_shape[-1::-1]/2.0 + img_shape[-1::-1]/2.0).astype(np.int32)
               print loc_this
               image_this = self.image[i,:,:,:]
               gtAlpha_this = self.gtAlpha[i,:,:,:]
               strokeMap_this = self.strokeMap[i,:,:,:]
               strokeMap_tmp = self.genStroke(strokeMap_this, loc_this, image_this, gtAlpha_this)
               strokeMap[i,:,:,:] = strokeMap_tmp
          
          return strokeMap
          
          
